package cn.sciento.boot.api.customize.agent;

import java.lang.instrument.Instrumentation;
import java.util.Arrays;

public class HotSwapAgent {
  private static Instrumentation _INST;
  
  private static final String CLASS_LOADER = "org.springframework.boot.loader.LaunchedURLClassLoader";
  
  private static final String SWAP_CLIENT = "cn.sciento.boot.api.customize.agent.HotSwapClient";
  
  private static final String SWAP_CLIENT_SETTER = "setInstrumentation";
  
  public static Instrumentation getInstrumentation() {
    return _INST;
  }
  
  public static void premain(String agentArgs, Instrumentation inst) throws Throwable {
    agentmain(agentArgs, inst);
  }
  
  public static void agentmain(String agentArgs, Instrumentation inst) throws Throwable {
    ClassLoader classLoader = getClassLoader(inst);
    Class<?> clientClass = classLoader.loadClass("cn.sciento.boot.api.customize.agent.HotSwapClient");
    clientClass.getMethod("setInstrumentation", new Class[] { Instrumentation.class }).invoke(null, inst);
    _INST = inst;
  }
  
  private static ClassLoader getClassLoader(Instrumentation inst) {
    Class[] loadedClasses = inst.getAllLoadedClasses();
    return Arrays.stream(loadedClasses)
      .filter(c -> (c.getClassLoader() != null && c.getClassLoader().getClass() != null))
      .filter(c -> "org.springframework.boot.loader.LaunchedURLClassLoader".equals(c.getClassLoader().getClass().getName()))
      .map(Class::getClassLoader)
      .findFirst()
      .orElse(Thread.currentThread().getContextClassLoader());
  }
}
